package cn.bugstack.openai.session;

import cn.bugstack.openai.executor.Executor;
import cn.bugstack.openai.executor.model.aliyun.AliModelExecutor;
import cn.bugstack.openai.executor.model.aliyun.config.AliModelConfig;
import cn.bugstack.openai.executor.model.baichuan.BaiChuanModelExecutor;
import cn.bugstack.openai.executor.model.baichuan.config.BaiChuanModelConfig;
import cn.bugstack.openai.executor.model.chatglm.ChatGLMModelExecutor;
import cn.bugstack.openai.executor.model.chatglm.config.ChatGLMConfig;
import cn.bugstack.openai.executor.model.chatgpt.ChatGPTModelExecutor;
import cn.bugstack.openai.executor.model.chatgpt.config.ChatGPTConfig;
import cn.bugstack.openai.executor.model.xunfei.XunFeiModelExecutor;
import cn.bugstack.openai.executor.model.xunfei.config.XunFeiConfig;
import cn.bugstack.openai.executor.parameter.CompletionRequest;
import lombok.*;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import okhttp3.logging.HttpLoggingInterceptor;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSources;

import java.util.HashMap;

/**
 * 配置文件
 *
 * @author 小傅哥，微信：fustack
 */
@Slf4j
@Data
public class Configuration {

    /**
     * 百川AI Config
     */
    private BaiChuanModelConfig baiChuanModelConfig;

    /**
     * 智谱Ai ChatGLM Config
     */
    private ChatGLMConfig chatGLMConfig;

    /**
     * OpenAi ChatGLM Config
     */
    private ChatGPTConfig chatGPTConfig;

    /**
     * 讯飞
     */
    private XunFeiConfig xunFeiConfig;

    /**
     * 阿里通义千问
     */
    private AliModelConfig aliModelConfig;

    /**
     * OkHttpClient
     */
    private OkHttpClient okHttpClient;

    private HashMap<String, Executor> executorGroup;

    public EventSource.Factory createRequestFactory() {
        return EventSources.createFactory(okHttpClient);
    }

    // OkHttp 配置信息
    private HttpLoggingInterceptor.Level level = HttpLoggingInterceptor.Level.HEADERS;
    private long connectTimeout = 4500;
    private long writeTimeout = 4500;
    private long readTimeout = 4500;

    // http keywords
    public static final String SSE_CONTENT_TYPE = "text/event-stream";
    public static final String DEFAULT_USER_AGENT = "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)";
    public static final String APPLICATION_JSON = "application/json";
    public static final String JSON_CONTENT_TYPE = APPLICATION_JSON + "; charset=utf-8";

    public HashMap<String, Executor> newExecutorGroup() {
        this.executorGroup = new HashMap<>();
        // 百川大模型 填充
        BaiChuanModelExecutor baicChuanModelExecutor = new BaiChuanModelExecutor(this);
        executorGroup.put(CompletionRequest.Model.BAICHUAN.getCode(), baicChuanModelExecutor);
        // ChatGLM 类型执行器填充
        Executor chatGLMModelExecutor = new ChatGLMModelExecutor(this);
        executorGroup.put(CompletionRequest.Model.CHATGLM_TURBO.getCode(), chatGLMModelExecutor);
        // ChatGPT 类型执行器填充
        Executor chatGPTModelExecutor = new ChatGPTModelExecutor(this);
        executorGroup.put(CompletionRequest.Model.GPT_3_5_TURBO.getCode(), chatGPTModelExecutor);
        executorGroup.put(CompletionRequest.Model.GPT_3_5_TURBO_1106.getCode(), chatGPTModelExecutor);
        executorGroup.put(CompletionRequest.Model.GPT_3_5_TURBO_16K.getCode(), chatGPTModelExecutor);
        executorGroup.put(CompletionRequest.Model.GPT_4.getCode(), chatGPTModelExecutor);
        executorGroup.put(CompletionRequest.Model.GPT_4_32K.getCode(), chatGPTModelExecutor);
        // XUNFEI
        Executor xunfeiModelExecutor = new XunFeiModelExecutor(this);
        executorGroup.put(CompletionRequest.Model.XUNFEI.getCode(), xunfeiModelExecutor);
        // 阿里通义千问
        Executor aliModelExecutor = new AliModelExecutor(this);
        executorGroup.put(CompletionRequest.Model.QWEN_TURBO.getCode(), aliModelExecutor);
        executorGroup.put(CompletionRequest.Model.QWEN_PLUS.getCode(), aliModelExecutor);
        executorGroup.put(CompletionRequest.Model.QWEN_MAX.getCode(), aliModelExecutor);

        return this.executorGroup;
    }

}
